import sys
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import sympy as sp
from sympy.utilities.lambdify import lambdify
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QLineEdit, QPushButton, QRadioButton, QButtonGroup, 
                            QCheckBox, QGroupBox, QGridLayout, QMessageBox, QComboBox)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont
from PyQt5.QtWebEngineWidgets import QWebEngineView

class FunctionPlotter(QMainWindow):
    """函数绘图器主窗口"""
    def __init__(self):
        super().__init__()
        self.setWindowTitle('函数图像绘制工具')
        self.setGeometry(100, 100, 1200, 800)
        
        # 初始化变量
        self.x = sp.Symbol('x')
        self.current_function = "x**2"
        self.current_x_value = 1
        self.current_range = [-10, 10]
        self.current_type = "function"
        self.plot_handles = []
        self.curve_params = {"a": 1, "b": 1, "c": 0, "curve_type": "ellipse"}
        self.default_curve_params = {
            "ellipse": {"a": 3, "b": 2, "c": 0},
            "hyperbola": {"a": 2, "b": 1, "c": 0},
            "parabola": {"a": 0.5, "b": 0, "c": 0}
        }
        self.show_grid = True
        self.show_legend = True
        self.show_axes = True
        self.show_ticks = True
        self.show_labels = True
        
        # 创建主布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)
        
        # 创建顶部控制区
        self.create_control_panel()
        
        # 创建绘图区域
        self.create_plot_area()
        
        # 初始不绘制图像，等待用户点击按钮
        # self.update_plot()  # 注释掉自动绘制
    
    def create_control_panel(self):
        """创建控制面板"""
        control_group = QGroupBox("绘图控制")
        control_group.setFont(QFont("SimHei", 10, QFont.Weight.Bold))
        control_layout = QGridLayout(control_group)
        
        # 函数输入
        control_layout.addWidget(QLabel("函数 f(x):"), 0, 0)
        self.func_input = QLineEdit(self.current_function)
        self.func_input.returnPressed.connect(self.update_plot)
        control_layout.addWidget(self.func_input, 0, 1)
        
        # X值输入（用于切线）
        control_layout.addWidget(QLabel("X值:"), 0, 2)
        self.x_value_input = QLineEdit(str(self.current_x_value))
        self.x_value_input.returnPressed.connect(self.update_plot)
        control_layout.addWidget(self.x_value_input, 0, 3)
        
        # 范围输入
        control_layout.addWidget(QLabel("范围:"), 0, 4)
        self.range_input = QLineEdit(str(self.current_range))
        self.range_input.returnPressed.connect(self.update_plot)
        control_layout.addWidget(self.range_input, 0, 5)
        
        # 绘图类型选择
        type_group = QGroupBox("绘图类型")
        type_layout = QVBoxLayout(type_group)
        
        self.plot_type_buttons = QButtonGroup()
        plot_types = ["函数", "导函数", "切线", "圆锥曲线"]
        for i, plot_type in enumerate(plot_types):
            radio = QRadioButton(plot_type)
            type_layout.addWidget(radio)
            self.plot_type_buttons.addButton(radio, i)
            if plot_type == "函数":
                radio.setChecked(True)
        
        self.plot_type_buttons.buttonClicked.connect(self.update_plot)
        control_layout.addWidget(type_group, 1, 0, 2, 2)
        
        # 圆锥曲线参数
        control_layout.addWidget(QLabel("参数 a:"), 1, 2)
        self.param_a_input = QLineEdit(str(self.curve_params["a"]))
        self.param_a_input.returnPressed.connect(self.update_plot)
        control_layout.addWidget(self.param_a_input, 1, 3)
        
        control_layout.addWidget(QLabel("参数 b:"), 1, 4)
        self.param_b_input = QLineEdit(str(self.curve_params["b"]))
        self.param_b_input.returnPressed.connect(self.update_plot)
        control_layout.addWidget(self.param_b_input, 1, 5)
        
        control_layout.addWidget(QLabel("参数 c:"), 2, 2)
        self.param_c_input = QLineEdit(str(self.curve_params["c"]))
        self.param_c_input.returnPressed.connect(self.update_plot)
        control_layout.addWidget(self.param_c_input, 2, 3)
        
        # 圆锥曲线类型选择
        control_layout.addWidget(QLabel("圆锥曲线类型:"), 2, 4)
        self.curve_type_combo = QComboBox()
        self.curve_type_combo.addItems(["椭圆", "双曲线", "抛物线"])
        self.curve_type_combo.currentTextChanged.connect(self.on_curve_type_changed)
        control_layout.addWidget(self.curve_type_combo, 2, 5)
        
        # 绘制图像按钮
        self.plot_button = QPushButton("绘制图像")
        self.plot_button.setStyleSheet("background-color: #007BFF; color: white; font-size: 14px;")
        self.plot_button.setMinimumHeight(40)
        self.plot_button.clicked.connect(self.update_plot)
        control_layout.addWidget(self.plot_button, 3, 0, 1, 2)
        
        # 刷新和清除按钮
        button_layout = QHBoxLayout()
        
        self.refresh_button = QPushButton("刷新")
        self.refresh_button.setStyleSheet("background-color: #28a745; color: white;")
        self.refresh_button.clicked.connect(self.update_plot)
        button_layout.addWidget(self.refresh_button)
        
        self.clear_button = QPushButton("清除")
        self.clear_button.setStyleSheet("background-color: #dc3545; color: white;")
        self.clear_button.clicked.connect(self.clear_plot)
        button_layout.addWidget(self.clear_button)
        
        control_layout.addLayout(button_layout, 3, 2, 1, 4)
        
        # 绘图选项
        options_layout = QGridLayout()
        
        self.grid_check = QCheckBox("网格线")
        self.grid_check.setChecked(self.show_grid)
        self.grid_check.stateChanged.connect(self.on_grid_changed)
        options_layout.addWidget(self.grid_check, 0, 0)
        
        self.legend_check = QCheckBox("图例")
        self.legend_check.setChecked(self.show_legend)
        self.legend_check.stateChanged.connect(self.on_legend_changed)
        options_layout.addWidget(self.legend_check, 0, 1)
        
        self.axes_check = QCheckBox("坐标轴")
        self.axes_check.setChecked(self.show_axes)
        self.axes_check.stateChanged.connect(self.on_axes_changed)
        options_layout.addWidget(self.axes_check, 0, 2)
        
        self.ticks_check = QCheckBox("刻度")
        self.ticks_check.setChecked(self.show_ticks)
        self.ticks_check.stateChanged.connect(self.on_ticks_changed)
        options_layout.addWidget(self.ticks_check, 0, 3)
        
        self.labels_check = QCheckBox("标签")
        self.labels_check.setChecked(self.show_labels)
        self.labels_check.stateChanged.connect(self.on_labels_changed)
        options_layout.addWidget(self.labels_check, 0, 4)
        
        control_layout.addLayout(options_layout, 4, 0, 1, 6)
        
        self.main_layout.addWidget(control_group)
    
    def create_plot_area(self):
        """创建绘图区域"""
        plot_group = QGroupBox("函数图像")
        plot_group.setFont(QFont("SimHei", 10, QFont.Weight.Bold))
        plot_layout = QVBoxLayout(plot_group)
        
        # 创建WebEngineView用于显示Plotly图表
        self.web_view = QWebEngineView()
        plot_layout.addWidget(self.web_view)
        
        self.main_layout.addWidget(plot_group)
    
    def on_curve_type_changed(self, text):
        """当圆锥曲线类型改变时，设置默认参数"""
        curve_types = {"椭圆": "ellipse", "双曲线": "hyperbola", "抛物线": "parabola"}
        curve_type = curve_types[text]
        
        # 更新当前参数为默认值
        self.curve_params = self.default_curve_params[curve_type].copy()
        self.curve_params["curve_type"] = curve_type
        
        # 更新UI显示
        self.param_a_input.setText(str(self.curve_params["a"]))
        self.param_b_input.setText(str(self.curve_params["b"]))
        self.param_c_input.setText(str(self.curve_params["c"]))
        
        # 如果当前是圆锥曲线类型，自动更新绘图
        if self.current_type == "conic":
            self.update_plot()
    
    def on_grid_changed(self, state):
        self.show_grid = state == Qt.Checked
        self.update_plot()
    
    def on_legend_changed(self, state):
        self.show_legend = state == Qt.Checked
        self.update_plot()
    
    def on_axes_changed(self, state):
        self.show_axes = state == Qt.Checked
        self.update_plot()
    
    def on_ticks_changed(self, state):
        self.show_ticks = state == Qt.Checked
        self.update_plot()
    
    def on_labels_changed(self, state):
        self.show_labels = state == Qt.Checked
        self.update_plot()
    
    def update_function(self):
        """更新函数表达式"""
        self.current_function = self.func_input.text()
    
    def update_x_value(self):
        """更新X值（用于切线）"""
        try:
            self.current_x_value = float(self.x_value_input.text())
        except ValueError:
            self.show_error("无效的X值，请输入一个有效的数字")
    
    def update_range(self):
        """更新绘图范围"""
        try:
            # 解析输入的范围文本，例如: [-5, 5]
            range_str = self.range_input.text().strip('[]').split(',')
            self.current_range = [float(range_str[0]), float(range_str[1])]
        except:
            self.show_error("无效的范围格式，请使用类似 [-5, 5] 的格式")
    
    def update_type(self):
        """更新绘图类型"""
        types = {0: 'function', 1: 'derivative', 2: 'tangent', 3: 'conic'}
        self.current_type = types[self.plot_type_buttons.checkedId()]
    
    def update_curve_params(self):
        """更新圆锥曲线参数"""
        try:
            self.curve_params["a"] = float(self.param_a_input.text())
            self.curve_params["b"] = float(self.param_b_input.text())
            self.curve_params["c"] = float(self.param_c_input.text())
        except ValueError as e:
            self.show_error(f"参数错误: {str(e)}")
    
    def toggle_grid(self):
        """切换网格线显示"""
        self.update_plot()
    
    def toggle_legend(self):
        """切换图例显示"""
        self.update_plot()
    
    def clear_plot(self):
        """清除当前绘图"""
        # 创建一个空图表
        fig = go.Figure()
        self.setup_cartesian_layout(fig)
        self.update_web_view(fig)
    
    def show_error(self, message):
        """显示错误消息"""
        QMessageBox.critical(self, "错误", message)
    
    def update_web_view(self, fig):
        """将Plotly图表更新到WebView中"""
        # 转换为HTML并加载
        html_content = fig.to_html(full_html=False, include_plotlyjs='cdn')
        self.web_view.setHtml(html_content)
    
    def setup_cartesian_layout(self, fig):
        """设置笛卡尔直角坐标系"""
        # 设置坐标轴范围
        x_min, x_max = self.current_range
        y_min, y_max = self.current_range
        
        # 设置布局
        fig.update_layout(
            title="函数图像绘制",
            font=dict(family="SimHei", size=12),
            showlegend=self.show_legend,
            template="seaborn",
            margin=dict(l=60, r=60, t=60, b=60),
            xaxis=dict(
                range=[x_min, x_max],
                showgrid=self.show_grid,
                gridcolor="rgba(200, 200, 200, 0.5)",
                zeroline=self.show_axes,
                zerolinecolor="black",
                zerolinewidth=2,
                showticklabels=self.show_ticks,
                title_text="x" if self.show_labels else "",
                title_font=dict(size=14)
            ),
            yaxis=dict(
                range=[y_min, y_max],
                showgrid=self.show_grid,
                gridcolor="rgba(200, 200, 200, 0.5)",
                zeroline=self.show_axes,
                zerolinecolor="black",
                zerolinewidth=2,
                showticklabels=self.show_ticks,
                title_text="y" if self.show_labels else "",
                title_font=dict(size=14)
            )
        )
    
    def update_plot(self):
        """更新绘图"""
        # 先更新所有参数
        self.update_function()
        self.update_x_value()
        self.update_range()
        self.update_type()
        if self.current_type == 'conic':
            self.update_curve_params()
        
        try:
            # 创建图表
            fig = make_subplots(rows=1, cols=1)
            
            if self.current_type != 'conic':
                # 解析函数表达式
                f_expr = sp.sympify(self.current_function)
                f = lambdify(self.x, f_expr, 'numpy')
                
                # 生成x值
                x_vals = np.linspace(self.current_range[0], self.current_range[1], 1000)
                
                if self.current_type == 'function':
                    # 绘制函数图像
                    y_vals = f(x_vals)
                    fig.add_trace(
                        go.Scatter(
                            x=x_vals, 
                            y=y_vals, 
                            mode='lines',
                            name=f'f(x) = {self.current_function}',
                            line=dict(color='blue', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    fig.update_layout(title=f'函数图像: f(x) = {self.current_function}')
                
                elif self.current_type == 'derivative':
                    # 计算导函数
                    f_prime_expr = sp.diff(f_expr, self.x)
                    f_prime = lambdify(self.x, f_prime_expr, 'numpy')
                    
                    # 绘制导函数图像
                    y_vals = f_prime(x_vals)
                    fig.add_trace(
                        go.Scatter(
                            x=x_vals, 
                            y=y_vals, 
                            mode='lines',
                            name=f"f'(x) = {f_prime_expr}",
                            line=dict(color='red', width=2),
                            hovertemplate='x: %{x:.2f}<br>f\'(x): %{y:.2f}<extra></extra>'
                        )
                    )
                    fig.update_layout(title=f"导函数图像: f'(x) = {f_prime_expr}")
                
                elif self.current_type == 'tangent':
                    # 计算切线
                    f_prime_expr = sp.diff(f_expr, self.x)
                    f_prime = lambdify(self.x, f_prime_expr, 'numpy')
                    
                    # 计算给定点的函数值和导数值
                    x0 = self.current_x_value
                    y0 = f(x0)
                    slope = f_prime(x0)
                    
                    # 切线方程: y = slope * (x - x0) + y0
                    tangent_line = lambda x: slope * (x - x0) + y0
                    
                    # 绘制原函数
                    fig.add_trace(
                        go.Scatter(
                            x=x_vals, 
                            y=f(x_vals), 
                            mode='lines',
                            name=f'f(x) = {self.current_function}',
                            line=dict(color='blue', width=2),
                            hovertemplate='x: %{x:.2f}<br>f(x): %{y:.2f}<extra></extra>'
                        )
                    )
                    
                    # 绘制切线
                    fig.add_trace(
                        go.Scatter(
                            x=x_vals, 
                            y=tangent_line(x_vals), 
                            mode='lines',
                            name=f'切线 at x={x0}',
                            line=dict(color='red', width=1.5, dash='dash'),
                            hovertemplate='x: %{x:.2f}<br>切线: %{y:.2f}<extra></extra>'
                        )
                    )
                    
                    # 绘制切点
                    fig.add_trace(
                        go.Scatter(
                            x=[x0], 
                            y=[y0], 
                            mode='markers',
                            name=f'切点 ({x0}, {y0:.2f})',
                            marker=dict(color='green', size=8),
                            hovertemplate='切点: (%{x:.2f}, %{y:.2f})<extra></extra>'
                        )
                    )
                    
                    fig.update_layout(title=f'函数及其在 x={x0} 处的切线')
            
            else:  # 圆锥曲线
                # 解析圆锥曲线参数
                a = self.curve_params["a"]
                b = self.curve_params["b"]
                c = self.curve_params["c"]
                curve_type = self.curve_params["curve_type"]
                
                if curve_type == 'ellipse':
                    # 椭圆参数方程: (x-h)²/a² + (y-k)²/b² = 1
                    t = np.linspace(0, 2*np.pi, 1000)
                    x_ellipse = a * np.cos(t)
                    y_ellipse = b * np.sin(t)
                    
                    fig.add_trace(
                        go.Scatter(
                            x=x_ellipse, 
                            y=y_ellipse, 
                            mode='lines',
                            name=f'椭圆: x²/{a}² + y²/{b}² = 1',
                            line=dict(color='green', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    fig.update_layout(title=f'椭圆: x²/{a}² + y²/{b}² = 1')
                
                elif curve_type == 'hyperbola':
                    # 双曲线参数方程: x²/a² - y²/b² = 1
                    x_min, x_max = self.current_range
                    x_range = np.linspace(x_min, x_max, 1000)
                    
                    # 确保在双曲线定义域内
                    x_pos = x_range[abs(x_range) >= max(a, 0.1)]  # 避免除以零
                    
                    # 计算上下分支
                    y_upper = b * np.sqrt((x_pos**2/a**2) - 1)
                    y_lower = -b * np.sqrt((x_pos**2/a**2) - 1)
                    
                    # 绘制右支
                    fig.add_trace(
                        go.Scatter(
                            x=x_pos[x_pos > 0],
                            y=y_upper[x_pos > 0],
                            mode='lines',
                            name=f'双曲线: x²/{a}² - y²/{b}² = 1',
                            line=dict(color='red', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    
                    # 绘制右支下半部分
                    fig.add_trace(
                        go.Scatter(
                            x=x_pos[x_pos > 0],
                            y=y_lower[x_pos > 0],
                            mode='lines',
                            line=dict(color='red', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    
                    # 绘制左支
                    fig.add_trace(
                        go.Scatter(
                            x=x_pos[x_pos < 0],
                            y=y_upper[x_pos < 0],
                            mode='lines',
                            line=dict(color='red', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    
                    # 绘制左支下半部分
                    fig.add_trace(
                        go.Scatter(
                            x=x_pos[x_pos < 0],
                            y=y_lower[x_pos < 0],
                            mode='lines',
                            line=dict(color='red', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    
                    fig.update_layout(title=f'双曲线: x²/{a}² - y²/{b}² = 1')
                
                elif curve_type == 'parabola':
                    # 抛物线方程: y = ax² + bx + c
                    x_parabola = np.linspace(self.current_range[0], self.current_range[1], 1000)
                    y_parabola = a * x_parabola ** 2 + b * x_parabola + c
                    
                    fig.add_trace(
                        go.Scatter(
                            x=x_parabola, 
                            y=y_parabola, 
                            mode='lines',
                            name=f'抛物线: y = {a}x² + {b}x + {c}',
                            line=dict(color='blue', width=2),
                            hovertemplate='x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>'
                        )
                    )
                    fig.update_layout(title=f'抛物线: y = {a}x² + {b}x + {c}')
            
            # 设置坐标轴比例
            if curve_type in ['ellipse', 'hyperbola']:
                fig.update_layout(
                    yaxis=dict(
                        scaleanchor="x",
                        scaleratio=1,
                    )
                )
            
            # 设置笛卡尔坐标系
            self.setup_cartesian_layout(fig)
            
            # 更新WebView
            self.update_web_view(fig)
            
        except Exception as e:
            self.show_error(f"绘图错误: {str(e)}")

if __name__ == "__main__":
    # 创建并运行应用程序
    app = QApplication(sys.argv)
    window = FunctionPlotter()
    window.show()
    sys.exit(app.exec_())    